package com.allm.ai.common.strategy.impl;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import org.springframework.stereotype.Component;

import com.allm.ai.common.dto.ApiRequest;
import com.allm.ai.common.dto.StreamingOutput;
import com.allm.ai.common.entity.AiModel;
import com.allm.ai.common.strategy.BaseAiModelStrategy;
import com.google.genai.Client;
import com.google.genai.ResponseStream;
import com.google.genai.types.Content;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.GenerateContentResponse;
import com.google.genai.types.Part;
import com.google.genai.types.ThinkingConfig;

import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;

@Slf4j
@Component
public class GoogleAIStrategy extends BaseAiModelStrategy {
    
    @Override
    public String callApi(AiModel model, ApiRequest request) throws Exception {
        log.info("调用 Google AI API，模型: {}", model.getName());
        
        try {
            // 创建 Google AI 客户端
            Client client = Client.builder().vertexAI(true)
                .project("lxdproject-464310")
                .location("us-west1")
                // .httpOptions(HttpOptions.builder().apiVersion("v1alpha").build())
                // .clientOptions(ClientOptions.builder().maxConnections(64).maxConnectionsPerHost(16).build())
                // .apiKey(model.getApiKey())
                .build();
            
            // 构建请求内容
            List<Content> contents = buildContents(request.getMessages());
            
            // 构建生成配置
            GenerateContentConfig config = GenerateContentConfig.builder()
                .temperature(request.getTemperature().floatValue())
                .maxOutputTokens(request.getMaxTokens())
                .topP(0.8f)
                .topK(40f)
                .build();
            
            // 发送请求
            GenerateContentResponse response = client.models.generateContent(
                model.getRequestModelName(), 
                contents, 
                config
            );
            
            // 解析响应
            return parseResponse(response);
            
        } catch (Exception e) {
            log.error("调用 Google AI API 失败", e);
            throw new RuntimeException("Google AI API 调用失败: " + e.getMessage(), e);
        }
    }
    
    /**
     * 构建内容列表
     */
    private List<Content> buildContents(List<Map<String, Object>> messages) {
        List<Content> contents = new ArrayList<>();
        
        for (Map<String, Object> message : messages) {
            String role = convertRole(message.get("role"));
            String content = String.valueOf(message.get("content"));
            
            Content.Builder contentBuilder = Content.builder()
                .role(role)
                .parts(List.of(Part.builder().text(content).build()));
            
            contents.add(contentBuilder.build());
        }
        
        return contents;
    }
    
    /**
     * 转换角色格式
     */
    private String convertRole(Object role) {
        String roleStr = String.valueOf(role);
        switch (roleStr) {
            case "user":
                return "user";
            case "assistant":
                return "model";
            case "system":
                return "user"; // Google AI 没有 system 角色，转换为 user
            default:
                return "user";
        }
    }
    
    /**
     * 解析响应
     */
    private String parseResponse(GenerateContentResponse response) {
        try {
            // 检查响应是否为空
            if (response == null) {
                throw new RuntimeException("Google AI API 响应为 null");
            }
            
            // 检查是否有候选响应
            if (response.candidates() == null || response.candidates().isEmpty()) {
                // 检查是否有错误信息
                if (response.promptFeedback() != null && response.promptFeedback().isPresent() && response.promptFeedback().get().blockReason() != null) {
                    throw new RuntimeException("请求被阻止: " + response.promptFeedback().get().blockReason());
                }
                throw new RuntimeException("Google AI API 响应为空，没有候选内容");
            }
            
            var candidate = response.candidates().get().get(0);
            
            // 检查候选内容是否被过滤
            if (candidate.finishReason() != null) {
                String finishReason = candidate.finishReason().toString();
                if ("SAFETY".equals(finishReason) || "RECITATION".equals(finishReason)) {
                    throw new RuntimeException("内容被安全过滤器阻止，原因: " + finishReason);
                }
            }
            
            // 检查内容
            if (candidate.content() == null || !candidate.content().isPresent()) {
                throw new RuntimeException("候选内容为空");
            }
            
            if (candidate.content().get().parts() == null || candidate.content().get().parts().isEmpty()) {
                throw new RuntimeException("候选内容部分为空");
            }
            
            // 拼接所有文本内容
            StringBuilder result = new StringBuilder();
            candidate.content().get().parts().get().forEach(part -> {
                if (part.text() != null && !part.text().get().trim().isEmpty()) {
                    result.append(part.text());
                }
            });
            
            String finalResult = result.toString().trim();
            if (finalResult.isEmpty()) {
                throw new RuntimeException("解析后的响应内容为空");
            }
            
            return finalResult;
            
        } catch (Exception e) {
            log.error("解析 Google AI 响应失败", e);
            throw new RuntimeException("解析 Google AI 响应失败: " + e.getMessage(), e);
        }
    }
    
    @Override
    public boolean supports(String serviceProvider) {
        return "google".equalsIgnoreCase(serviceProvider) || 
               "google-ai".equalsIgnoreCase(serviceProvider) ||
               "gemini".equalsIgnoreCase(serviceProvider);
    }
    
    /**
     * Google AI SDK 1.15.0 支持流式调用
     */
    @Override
    public boolean supportsStreaming() {
        return true;
    }
    
    /**
     * 流式调用实现
     */
    @Override
    public Flux<StreamingOutput> callApiStream(AiModel model, ApiRequest request) {
        return Flux.create(sink -> {
            try {
                log.info("调用 Google AI 流式 API，模型: {}", model.getName());
                
                // 创建客户端
                Client client = Client.builder()
                    .apiKey(model.getApiKey())
                    .build();
                
                // 构建请求内容
                List<Content> contents = buildContents(request.getMessages());
                
                // 构建生成配置
                GenerateContentConfig config = GenerateContentConfig.builder()
                    .temperature(request.getTemperature().floatValue())
                    .maxOutputTokens(request.getMaxTokens())
                    .thinkingConfig(ThinkingConfig.builder()
                        .thinkingBudget(-1)
                        .includeThoughts(true).build())
                    .topP(0.8f)
                    .topK(40f)
                    .build();

                // ResponseStream<GenerateContentResponse> responseStream =
                //     client.models.generateContentStream(
                //         model.getRequestModelName(),
                //             contents,
                //             null);
                // StringBuilder response = new StringBuilder();
                // for (GenerateContentResponse res : responseStream) {
                //     sink.next(new StreamingOutput(StreamingOutput.CONTENT, res.text()));
                //     // System.out.print(res.text());
                //     // response.append(res.text());
                // }
                // // sink.next(new StreamingOutput(StreamingOutput.CONTENT, response.toString()));
                // responseStream.close();
                
                // 流式调用
                client.models.generateContentStream(model.getRequestModelName(), contents, config)
                    .forEach(response -> {
                        try {
                            processStreamResponse(response, sink);
                        } catch (Exception e) {
                            log.error("处理流式响应失败", e);
                            sink.error(e);
                        }
                    });
                
                sink.complete();

                // CompletableFuture<ResponseStream<GenerateContentResponse>> responseFuture =
                //     client.async.models.generateContentStream(model.getName(), contents, config);

                // responseFuture.thenAccept(
                //     response -> {
                //         for (GenerateContentResponse res : responseStream) {
                //             sink.next(new StreamingOutput(StreamingOutput.CONTENT, res.text()));
                //             System.out.println("Async response: " + res.text());
                //             // System.out.print(res.text());
                //             // response.append(res.text());
                //         }
                //     })
                // .join();
                
            } catch (Exception e) {
                log.error("Google AI 流式调用失败", e);
                sink.error(e);
            }
        });
    }
    
    /**
     * 处理流式响应 - 使用父类的分块能力
     */
    private void processStreamResponse(GenerateContentResponse response, reactor.core.publisher.FluxSink<StreamingOutput> sink) {
        try {
            // 检查响应是否为空
            if (response == null || response.candidates() == null || response.candidates().isEmpty()) {
                return;
            }
            
            var candidate = response.candidates().get().get(0);
            
            // 检查候选内容
            if (candidate.content() == null || !candidate.content().isPresent()) {
                return;
            }
            
            var content = candidate.content().get();
            if (content.parts() == null) {
                return;
            }
            
            if (content.parts() == null || !content.parts().isPresent() || content.parts().get() == null) {
                return;
            }
            
            List<Part> parts = content.parts().get();
            parts.forEach(part -> {
                if (part.text() != null && !part.text().get().trim().isEmpty()) {
                    String text = part.text().get();
                    
                    // 使用父类的分块输出能力
                    extractThinkingContentWithChunking(text, sink, true); // true表示使用分块输出
                }
            });
            
        } catch (Exception e) {
            log.error("处理流式响应失败", e);
            sink.error(e);
        }
    }
} 